/**
 * Copyright (C) 2010 - 2013 Harry Glasgow
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package com.googlecode.jaden.ui.dialog.manualtraining;

import com.googlecode.jaden.common.config.IOConfiguration;
import com.googlecode.jaden.common.enums.IOType;
import com.googlecode.jaden.common.other.Utils;
import com.googlecode.jaden.common.enums.LayerType;
import com.googlecode.jaden.engine.Layer;
import static com.googlecode.jaden.ui.Constants.*;
import com.googlecode.jaden.ui.UI;
import com.googlecode.jaden.ui.UiUtils;
import com.googlecode.jaden.ui.dialog.BaseDialog;

import javax.swing.*;
import java.awt.*;
import java.awt.datatransfer.Clipboard;
import java.awt.datatransfer.ClipboardOwner;
import java.awt.datatransfer.StringSelection;
import java.awt.datatransfer.Transferable;
import java.awt.event.ActionEvent;
import java.awt.event.MouseAdapter;
import java.awt.event.MouseEvent;
import java.util.prefs.Preferences;
import java.util.Map;

public class AnalysisDialog extends BaseDialog implements ClipboardOwner {

    private final IOConfiguration[] inputConfigurations;
    private final Map<Integer, Layer> layers;
    private final IOConfiguration outputConfiguration;

    private JTextArea textArea;

    private static final int SPACER = 10;
    private static final int TEXT_OFFSET = 15;

    public AnalysisDialog(UI ui, Preferences preferences, IOConfiguration[] inputConfigurations,
                          Map<Integer, Layer> layers, IOConfiguration outputConfiguration) {
        super(ui, DIALOG_NETWORK_ANALYSIS_TITLE);
        this.inputConfigurations = inputConfigurations;
        this.layers = layers;
        this.outputConfiguration = outputConfiguration;

        removeCancel();
        JTabbedPane tabbedPane = new JTabbedPane();
        addTextPane(tabbedPane);
        addGraphPane(tabbedPane);
        setMainPanel(UiUtils.verticalList(tabbedPane), preferences);
    }

    private void addTextPane(JTabbedPane tabbedPane) {
        textArea = new JTextArea(10, 40);
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < layers.size(); i++) {
            LayerType layerType = layers.get(i).getLayerType();
            double[][] weights = layers.get(i).getWeights();
            double[] biases = layers.get(i).getBiases();
            String layerText = Utils.translate(LABEL_LAYER_TEXT) + ' ' + (1 + i);
            sb.append(layerText + '\n');
            sb.append(Utils.underline(layerText, true) + '\n');
            for (int j = 0; j < biases.length; j++) {
                String neuronText = Utils.translate(LABEL_NEURON_TEXT) + ' ' + (j + 1);
                sb.append(neuronText + '\n');
                sb.append(Utils.underline(neuronText, false) + '\n');
                sb.append(Utils.translate(LABEL_LAYER_TYPE_TEXT) + ' ' +
                        layerType.getDescription() + '\n');
                sb.append(Utils.translate(LABEL_BIAS_TEXT) + ' ' + biases[j] + '\n');
                for (int k = 0; k < weights[j].length; k++) {
                    sb.append(Utils.translate(LABEL_WEIGHT_TEXT) + ' ' + (k + 1) + ' ' +
                            weights[j][k] + '\n');
                }
            }

            sb.append('\n');
        }
        textArea.setText(sb.toString());
        JScrollPane scrollPane = new JScrollPane(textArea);
        textArea.addMouseListener(new MouseAdapter() {
            public void mousePressed(MouseEvent e) {
                maybeShowPopup(e);
            }

            public void mouseReleased(MouseEvent e) {
                maybeShowPopup(e);
            }

            private void maybeShowPopup(MouseEvent e) {
                if (e.isPopupTrigger()) {
                    showPopup(e);
                }
            }
        });
        scrollPane.setPreferredSize(new Dimension(600, 375));
        tabbedPane.add(Utils.translate(LABEL_TEXT_TEXT), scrollPane);
    }

    private void showPopup(MouseEvent e) {
        JPopupMenu menu = new JPopupMenu();
        menu.add(new AbstractAction("Copy") {
            public void actionPerformed(ActionEvent e) {
                copyText();
            }
        });
        menu.add(new AbstractAction("Select All") {
            public void actionPerformed(ActionEvent e) {
                selectAllText();
            }
        });
        menu.show(e.getComponent(), e.getX(), e.getY());
    }

    private void copyText() {
        StringSelection stringSelection = new StringSelection(textArea.getSelectedText());
        Clipboard clipboard = Toolkit.getDefaultToolkit().getSystemClipboard();
        clipboard.setContents(stringSelection, this);
    }

    private void selectAllText() {
        textArea.requestFocus();
        textArea.setSelectionStart(0);
        textArea.setSelectionEnd(textArea.getText().length());
    }

    public void lostOwnership(Clipboard clipboard, Transferable contents) {
    }

    private void addGraphPane(JTabbedPane tabbedPane) {
        JPanel graphPanel = new JPanel() {

            public void paint(Graphics g) {

                g.setColor(Color.white);
                g.fillRect(0, 0, getPreferredSize().width, getPreferredSize().height);

                int panelHeight = getPreferredSize().height;

                // Input text
                g.setColor(Color.black);
                int inputsHeight = 0;
                for (IOConfiguration ioConfiguration : inputConfigurations) {
                    int inputs = ioConfiguration.getInternalWidth();
                    int inputHeight = SPACER + SPACER * inputs;
                    inputsHeight += inputHeight;
                }
                int cumulativeY = panelHeight / 2 - inputsHeight / 2;
                int cumulativeX = SPACER;
                for (IOConfiguration inputConfiguration : inputConfigurations) {
                    int inputs = inputConfiguration.getInternalWidth();
                    int inputHeight = SPACER + SPACER * inputs;
                    g.drawString(inputConfiguration.getName(), cumulativeX, cumulativeY + TEXT_OFFSET);
                    cumulativeY += inputHeight;
                }
                cumulativeX += calculateInputNameWidth() + SPACER;

                // Input lines
                cumulativeY = panelHeight / 2 - inputsHeight / 2;
                for (IOConfiguration ioConfiguration : inputConfigurations) {
                    int connections = ioConfiguration.getInternalWidth();
                    int inputHeight = SPACER + SPACER * connections;
                    // Left half of input.
                    g.setColor(Color.black);
                    g.drawLine(cumulativeX, SPACER + cumulativeY, cumulativeX + SPACER, SPACER + cumulativeY);
                    // Right, multiple half of input.
                    for (int connection = 0; connection < connections; connection++) {
                        // Connection line.
                        g.setColor(Color.black);
                        g.drawLine(SPACER + cumulativeX, SPACER + cumulativeY + SPACER * connection,
                                SPACER + cumulativeX + SPACER, SPACER + cumulativeY + SPACER * connection);
                        // Connection dot at right of input line.
                        g.setColor(Color.black);
                        g.drawOval(SPACER * 2 + cumulativeX - 2, SPACER + cumulativeY + SPACER * connection - 2, 4, 4);
                    }
                    // Vertical connection joiner bar.
                    g.setColor(Color.black);
                    g.drawLine(SPACER + cumulativeX, cumulativeY + SPACER, SPACER + cumulativeX,
                            cumulativeY + SPACER + SPACER * (connections - 1));
                    cumulativeY += inputHeight;
                }
                cumulativeX += 2 * SPACER;

                double maxBiasOrWeight = presetColorDepth(layers);

                // Layers
                for (int layerNumber = 0; layerNumber < layers.size(); layerNumber++) {

                    Layer layerCurrent = layers.get(layerNumber);

                    // Connections
                    if (layerNumber == 0) {

                        // Connect to inputs
                        int cumulativeYLeft = panelHeight / 2 - inputsHeight / 2;
                        int numberOfNeurons = layerCurrent.getLayerWidth();
                        int inputConfigurationIndex = 0;
                        for (IOConfiguration ioConfiguration : inputConfigurations) {
                            int connections = ioConfiguration.getInternalWidth();
                            for (int connection = 0; connection < connections; connection++) {
                                for (int neuronNumber = 0; neuronNumber < numberOfNeurons; neuronNumber++) {
                                    g.setColor(calculateColor(layerCurrent.getWeights()[neuronNumber][inputConfigurationIndex], maxBiasOrWeight));
                                    cumulativeY = panelHeight / 2 - 4 * numberOfNeurons * SPACER / 2 +
                                            4 * SPACER * neuronNumber + SPACER;
                                    g.drawLine(cumulativeX,
                                            SPACER + cumulativeYLeft + SPACER * connection,
                                            cumulativeX + SPACER * 10,
                                            SPACER + cumulativeY);
                                }
                                inputConfigurationIndex++;
                            }
                            int inputHeight = SPACER + SPACER * connections;
                            cumulativeYLeft += inputHeight;
                        }
                    } else {

                        // Connect to previous layer
                        int numberOfNeuronsPrevious = layers.get(layerNumber - 1).getLayerWidth();
                        int numberOfNeurons = layerCurrent.getLayerWidth();
                        int cumulativeYLeft = panelHeight / 2 - 4 * numberOfNeuronsPrevious * SPACER / 2 + SPACER;
                        for (int neuronNumberLeft = 0; neuronNumberLeft < numberOfNeuronsPrevious; neuronNumberLeft++) {
                            cumulativeY = panelHeight / 2 - 4 * numberOfNeurons * SPACER / 2 + SPACER;
                            for (int neuronNumber = 0; neuronNumber < numberOfNeurons; neuronNumber++) {
                                g.setColor(calculateColor(layerCurrent.getWeights()[neuronNumber][neuronNumberLeft], maxBiasOrWeight));
                                g.drawLine(cumulativeX,
                                        SPACER + cumulativeYLeft,
                                        cumulativeX + SPACER * 10,
                                        SPACER + cumulativeY);
                                cumulativeY += 4 * SPACER;
                            }
                            cumulativeYLeft += 4 * SPACER;
                        }
                    }

                    // Layer
                    int numberOfNeurons = layerCurrent.getLayerWidth();
                    cumulativeY = panelHeight / 2 - 4 * numberOfNeurons * SPACER / 2;
                    double[] biasesCurrent = layerCurrent.getBiases();
                    for (int j = 0; j < numberOfNeurons; j++) {
                        // Neuron box.
                        g.setColor(Color.black);
                        g.drawRect(cumulativeX + 10 * SPACER, cumulativeY + SPACER / 2, 2 * SPACER, 3 * SPACER);
                        // Dot at center of neuron box.
                        g.setColor(calculateColor(biasesCurrent[j], maxBiasOrWeight));
                        g.drawOval(cumulativeX + 11 * SPACER - 2, cumulativeY + SPACER * 2 - 2, 4, 4);
                        cumulativeY += 4 * SPACER;
                    }
                    cumulativeX += 12 * SPACER;
                }

                // Join up for categorical outputs.
                if (outputConfiguration.getType() == IOType.Categorical) {
                    int numberOfNeurons = layers.get(layers.size() - 1).getLayerWidth();
                    cumulativeY = panelHeight / 2 - 4 * numberOfNeurons * SPACER / 2 + SPACER + 10;
                    int initialY = cumulativeY;
                    for (int i = 0; i < numberOfNeurons; i++) {
                        g.setColor(Color.black);
                        g.drawLine(cumulativeX, cumulativeY, cumulativeX + SPACER, cumulativeY);
                        if (i == numberOfNeurons - 1) {
                            g.drawLine(cumulativeX + SPACER, initialY, cumulativeX + SPACER, cumulativeY);
                        }
                        cumulativeY += 4 * SPACER;
                    }
                    cumulativeX += SPACER;
                }

                // Output lines
                g.setColor(Color.black);
                int length = outputConfiguration.getType() == IOType.Categorical ? 1 : 2;
                g.drawLine(cumulativeX, getPreferredSize().height / 2, cumulativeX + length * SPACER,
                        getPreferredSize().height / 2);
                cumulativeX += length * SPACER;

                // Output text
                g.setColor(Color.black);
                g.drawString(outputConfiguration.getName(), cumulativeX + SPACER,
                        getPreferredSize().height / 2 + TEXT_OFFSET - SPACER);
            }

            public Dimension getPreferredSize() {

                // *** Height ***

                int preferredHeight = 0;
                for (IOConfiguration ioConfiguration : inputConfigurations) {
                    preferredHeight += SPACER + SPACER * ioConfiguration.getInternalWidth();
                }
                for (int i = 0; i < layers.size(); i++) {
                    preferredHeight = Math.max(preferredHeight, 4 * SPACER * layers.get(i).getLayerWidth());
                }

                // *** Width ***

                int preferredWidth = calculateInputNameWidth() + SPACER;

                // Input lines
                preferredWidth += 3 * SPACER;

                // Layers
                preferredWidth += 12 * SPACER * layers.size();

                // Output lines
                preferredWidth += 2 * SPACER;

                FontMetrics metrics = Toolkit.getDefaultToolkit().getFontMetrics(getFont());
                int outputNameWidth = metrics.charsWidth(outputConfiguration.getName().toCharArray(), 0,
                        outputConfiguration.getName().length());

                preferredWidth += SPACER + outputNameWidth;

                // Add a spacer to give the bottom/right a margin.
                return new Dimension(preferredWidth + SPACER, preferredHeight + SPACER);
            }

            private int calculateInputNameWidth() {
                int inputNameWidth = 0;
                FontMetrics metrics = Toolkit.getDefaultToolkit().getFontMetrics(getFont());
                for (IOConfiguration ioConfiguration : inputConfigurations) {
                    String name = ioConfiguration.getName();
                    int width = metrics.charsWidth(name.toCharArray(), 0, name.length());
                    inputNameWidth = Math.max(inputNameWidth, width);
                }
                return inputNameWidth;
            }

            private double presetColorDepth(Map<Integer, Layer> layers) {
                double min = Double.MAX_VALUE;
                double max = Double.MIN_VALUE;
                for (Layer layer : layers.values()) {
                    for (double bias : layer.getBiases()) {
                        if (bias < min) {
                            min = bias;
                        }
                        if (bias > max) {
                            max = bias;
                        }
                    }
                    for (double[] weights : layer.getWeights()) {
                        for (double weight : weights) {
                            if (weight < min) {
                                min = weight;
                            }
                            if (weight > max) {
                                max = weight;
                            }
                        }
                    }
                }
                return Math.max(max, -min);
            }

            private Color calculateColor(double target, double max) {
                double k = target / max;
                if (k > 0) {
                    return new Color(0, (int) (255 * k), 0);
                } else {
                    return new Color((int) (255 * -k), 0, 0);
                }
            }
        };

        JScrollPane scrollPane = new JScrollPane(graphPanel);
        scrollPane.setPreferredSize(new Dimension(400, 250));

        tabbedPane.add(Utils.translate(LABEL_GRAPH_TEXT), scrollPane);
    }

    protected void doOk() {
        dispose();
    }
}